package com.johnguo.beam_test.java;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.join.CoGbkResult;
import org.apache.beam.sdk.transforms.join.CoGroupByKey;
import org.apache.beam.sdk.transforms.join.KeyedPCollectionTuple;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.TupleTag;

public class JoinDemo {
	static final List<String> LINES1 = Arrays.asList("johnguo man hansame", "johnguo man hansame", "annasui woman c");
	static final List<String> LINES2 = Arrays.asList("johnguo 25");
	// static final List<String> sexInfoTag = Arrays.asList("name","sex");
	// static final List<String> scoreInfoTag = Arrays.asList("name","score");

	public static void main(String... args) {
		final TupleTag<List<String>> sexInfoTag = new TupleTag<List<String>>();
		final TupleTag<List<String>> scoreInfoTag = new TupleTag<List<String>>();

		PipelineOptions options = PipelineOptionsFactory.fromArgs(args).withValidation().create();
		Pipeline p = Pipeline.create(options);
		PCollection<String> reader1 = p.apply(Create.of(LINES1)).setCoder(StringUtf8Coder.of());
		PCollection<String> reader2 = p.apply(Create.of(LINES2)).setCoder(StringUtf8Coder.of());
		PCollection<KV<String, List<String>>> sexInfo = reader1
				.apply(ParDo.of(new DoFn<String, KV<String, List<String>>>() {
					private static final long serialVersionUID = 1L;

					@ProcessElement
					public void processElement(ProcessContext c) {
						String[] words = c.element().split(" ");
						List<String> hello = new ArrayList<String>();
						for (int n = 0; n < words.length; n++) {
							hello.add(words[n]);
						}
						c.output(KV.of(words[0], hello));
					}
				}));
		PCollection<KV<String, List<String>>> scoreInfo = reader2
				.apply(ParDo.of(new DoFn<String, KV<String, List<String>>>() {
	
					private static final long serialVersionUID = 1L;

					@ProcessElement
					public void processElement(ProcessContext c) {
						String[] words = c.element().split(" ");
						List<String> list = new ArrayList<String>();
						for (int n = 0; n < words.length; n++) {
							list.add(words[n]);
						}
						c.output(KV.of(words[0], list));
					}
				}));

		PCollection<KV<String, CoGbkResult>> kvpCollection = KeyedPCollectionTuple.of(sexInfoTag, sexInfo)
				.and(scoreInfoTag, scoreInfo).apply(CoGroupByKey.<String>create());
		kvpCollection.apply(ParDo.of(new DoFn<KV<String, CoGbkResult>, String>() {
			private static final long serialVersionUID = 1L;

			@ProcessElement
			public void processElement(ProcessContext c) {
				c.element().getValue();
				System.out.println(c.element());
			}
		}));
		p.run().waitUntilFinish();
	}
}
